#include "filters/simple_filter_engine.h"
#include "utils/logger.h"

#include <algorithm>

SimpleFilterEngine::SimpleFilterEngine() {
    LOG_DEBUG("SimpleFilterEngine initialized.");
}

bool SimpleFilterEngine::addRule(const FilterRule& rule) {
    std::unique_lock lock(m_rulesMutex);

    /// Check if the rule ID already exists.
    if (rule.id != 0 && m_rules.find(rule.id) != m_rules.end()) {
        LOG_ERROR("Rule with ID " << rule.id << " already exists");
        return false;
    }

    FilterRule newRule = rule;
    if (newRule.id == 0) {
        newRule.id = generateRuleId();
    }

    m_rules[newRule.id] = newRule;

    sortRulesByPriority();

    LOG_INFO("Added rule: " << newRule.name << " (ID: " << newRule.id << ")");
    return true;
}

bool SimpleFilterEngine::removeRule(uint32_t ruleId) {
    std::unique_lock lock(m_rulesMutex);

    auto it = m_rules.find(ruleId);
    if (it == m_rules.end()) {
        LOG_WARNING("Rule with ID " << ruleId << " not found");
        return false;
    }

    LOG_INFO("Removed rule: " << it->second.name << " (ID: " << it->second.id << ")");
    m_rules.erase(it);

    auto priority_it = std::find(m_rulePriorityOrder.begin(), m_rulePriorityOrder.end(), ruleId);
    if (priority_it != m_rulePriorityOrder.end()) {
        m_rulePriorityOrder.erase(priority_it);
    }

    return true;
}

bool SimpleFilterEngine::updateRule(const FilterRule& rule) {
    if (rule.id == 0) {
        LOG_ERROR("Cannot update rule with ID 0");
        return false;
    }

    std::unique_lock lock(m_rulesMutex);

    auto it = m_rules.find(rule.id);
    if (it == m_rules.end()) {
        LOG_ERROR("Rule with ID " << rule.id << " not found");
        return false;
    }

    it->second = rule;

    sortRulesByPriority();

    LOG_INFO("Update rule: " << rule.name << " (ID: " << rule.id << ")");
    return true;
}

void SimpleFilterEngine::setRules(const std::vector<FilterRule>& rules) {
    std::unique_lock lock(m_rulesMutex);

    m_rules.clear();
    m_rulePriorityOrder.clear();

    for (const auto& rule : rules) {
        FilterRule newRule = rule;
        if (newRule.id == 0) {
            newRule.id = generateRuleId();
        }
        m_rules[newRule.id] = newRule;
    }

    sortRulesByPriority();

    LOG_INFO("Set " << rules.size() << " new rules");
}

std::vector<FilterRule> SimpleFilterEngine::getRules() const {
    std::shared_lock lock(m_rulesMutex);

    std::vector<FilterRule> rules;
    rules.reserve(m_rules.size());

    for (const auto& [_, v] : m_rules) {
        rules.push_back(v);
    }

    return rules;
}

FilterResult SimpleFilterEngine::processPacket(const Packet& packet) {
    m_stats.total_packets++;

    if (!m_enabled) {
        m_stats.accepted_packets++;
        return FilterResult(FilterAction::ACCEPT);
    }

    std::shared_lock lock(m_rulesMutex);

    /// Filter the packet according to the rule priority order.
    for (uint32_t ruleId : m_rulePriorityOrder) {
        const auto& rule = m_rules[ruleId];

        if (rule.match(packet)) {
            /// Update statistics.
            {
                std::unique_lock lock(m_statsMutex);
                m_stats.filtered_packets++;
                m_stats.rule_hit_counts[ruleId]++;

                if (rule.action == FilterAction::ACCEPT) {
                    m_stats.accepted_packets++;
                } else {
                    m_stats.dropped_packets++;
                }
            }

            /// Atomic operation.
            const_cast<FilterRule&>(rule).hit_count++;

            LOG_DEBUG("Packet matched rule: " << rule.name << " (Action:"
                      << (rule.action == FilterAction::ACCEPT ? "ACCEPT" : "DROP") << ")");
            return FilterResult(rule.action, rule.id, rule.name);
        }
    }

    /// No rules, use default action.
    {
        std::unique_lock lock(m_statsMutex);
        if (m_defaultAction == FilterAction::ACCEPT) {
            m_stats.accepted_packets++;
        } else {
            m_stats.dropped_packets++;
            m_stats.filtered_packets++;
        }
    }

    return FilterResult(m_defaultAction);
}

FilterStats SimpleFilterEngine::getStatistics() const {
    std::shared_lock lock(m_statsMutex);
    return m_stats;
}

void SimpleFilterEngine::resetStatistics() {
    std::unique_lock lock(m_statsMutex);
    m_stats.reset();

    std::shared_lock rules_lock(m_rulesMutex);
    for (auto& [_, v] : m_rules) {
        v.hit_count = 0;
    }

    LOG_DEBUG("Filter statistics reset");
}

void SimpleFilterEngine::sortRulesByPriority() {
    m_rulePriorityOrder.clear();
    m_rulePriorityOrder.reserve(m_rules.size());

    for (const auto& [k, _] : m_rules) {
        m_rulePriorityOrder.push_back(k);
    }

    std::sort(m_rulePriorityOrder.begin(), m_rulePriorityOrder.end(),
              [this](uint32_t a, uint32_t b) {
        return m_rules[a].priority < m_rules[b].priority;
    });
}

uint32_t SimpleFilterEngine::generateRuleId() {
    return m_nextRuleId++;
}